//! SSD1306 OLED Display Driver
//! 
//! This module provides a driver for SSD1306 OLED displays over I2C.
//! 
//! # Supported Displays
//! - 128x64 pixels
//! - 128x32 pixels
//! - Monochrome (1-bit per pixel)
//! 
//! # Typical I2C Address
//! - 0x3C (default)
//! - 0x3D (alternate)
//! 
//! # Hardware Connection
//! ```text
//! SSD1306    ->  RK3588
//! VCC        ->  3.3V
//! GND        ->  GND
//! SCL        ->  I2C SCL (e.g., GPIO4_B5 for I2C4)
//! SDA        ->  I2C SDA (e.g., GPIO4_B6 for I2C4)
//! ```

use crate::rk3588_i2c::{Rk3588I2c, I2cError};

/// SSD1306 Commands
const CMD_SET_CONTRAST: u8 = 0x81;
const CMD_ENTIRE_DISPLAY_ON_RESUME: u8 = 0xA4;
const CMD_ENTIRE_DISPLAY_ON: u8 = 0xA5;
const CMD_NORMAL_DISPLAY: u8 = 0xA6;
const CMD_INVERSE_DISPLAY: u8 = 0xA7;
const CMD_DISPLAY_OFF: u8 = 0xAE;
const CMD_DISPLAY_ON: u8 = 0xAF;

const CMD_SET_MEMORY_ADDR_MODE: u8 = 0x20;
const CMD_SET_COLUMN_ADDR: u8 = 0x21;
const CMD_SET_PAGE_ADDR: u8 = 0x22;

const CMD_SET_START_LINE: u8 = 0x40;
const CMD_SET_SEGMENT_REMAP: u8 = 0xA0;
const CMD_SET_MULTIPLEX_RATIO: u8 = 0xA8;
const CMD_SET_COM_SCAN_INC: u8 = 0xC0;
const CMD_SET_COM_SCAN_DEC: u8 = 0xC8;
const CMD_SET_DISPLAY_OFFSET: u8 = 0xD3;
const CMD_SET_COM_PINS: u8 = 0xDA;

const CMD_SET_DISPLAY_CLOCK_DIV: u8 = 0xD5;
const CMD_SET_PRECHARGE: u8 = 0xD9;
const CMD_SET_VCOM_DESELECT: u8 = 0xDB;
const CMD_CHARGE_PUMP: u8 = 0x8D;

/// Display Size Configuration
#[derive(Debug, Clone, Copy)]
pub enum DisplaySize {
    Size128x64,
    Size128x32,
}

impl DisplaySize {
    fn width(&self) -> u8 {
        128
    }

    fn height(&self) -> u8 {
        match self {
            DisplaySize::Size128x64 => 64,
            DisplaySize::Size128x32 => 32,
        }
    }

    fn pages(&self) -> u8 {
        self.height() / 8
    }
}

/// SSD1306 OLED Display Driver
pub struct Ssd1306<'a> {
    i2c: &'a mut Rk3588I2c,
    addr: u8,
    size: DisplaySize,
    buffer: [u8; 1024], // Max buffer for 128x64 display
}

impl<'a> Ssd1306<'a> {
    /// Create a new SSD1306 driver instance
    pub fn new(i2c: &'a mut Rk3588I2c, addr: u8, size: DisplaySize) -> Self {
        Self {
            i2c,
            addr,
            size,
            buffer: [0; 1024],
        }
    }

    /// Initialize the display
    pub fn init(&mut self) -> Result<(), I2cError> {
        // Display off
        self.write_cmd(CMD_DISPLAY_OFF)?;

        // Set display clock divide ratio/oscillator frequency
        self.write_cmd(CMD_SET_DISPLAY_CLOCK_DIV)?;
        self.write_cmd(0x80)?;

        // Set multiplex ratio
        self.write_cmd(CMD_SET_MULTIPLEX_RATIO)?;
        self.write_cmd(self.size.height() - 1)?;

        // Set display offset
        self.write_cmd(CMD_SET_DISPLAY_OFFSET)?;
        self.write_cmd(0x00)?;

        // Set start line
        self.write_cmd(CMD_SET_START_LINE | 0x00)?;

        // Charge pump
        self.write_cmd(CMD_CHARGE_PUMP)?;
        self.write_cmd(0x14)?; // Enable charge pump

        // Set memory addressing mode to horizontal
        self.write_cmd(CMD_SET_MEMORY_ADDR_MODE)?;
        self.write_cmd(0x00)?; // Horizontal addressing mode

        // Set segment re-map (flip horizontally)
        self.write_cmd(CMD_SET_SEGMENT_REMAP | 0x01)?;

        // Set COM output scan direction (flip vertically)
        self.write_cmd(CMD_SET_COM_SCAN_DEC)?;

        // Set COM pins hardware configuration
        self.write_cmd(CMD_SET_COM_PINS)?;
        match self.size {
            DisplaySize::Size128x64 => self.write_cmd(0x12)?,
            DisplaySize::Size128x32 => self.write_cmd(0x02)?,
        }

        // Set contrast
        self.write_cmd(CMD_SET_CONTRAST)?;
        self.write_cmd(0xCF)?;

        // Set pre-charge period
        self.write_cmd(CMD_SET_PRECHARGE)?;
        self.write_cmd(0xF1)?;

        // Set VCOMH deselect level
        self.write_cmd(CMD_SET_VCOM_DESELECT)?;
        self.write_cmd(0x40)?;

        // Entire display on (resume to RAM content)
        self.write_cmd(CMD_ENTIRE_DISPLAY_ON_RESUME)?;

        // Set normal display (not inverted)
        self.write_cmd(CMD_NORMAL_DISPLAY)?;

        // Display on
        self.write_cmd(CMD_DISPLAY_ON)?;

        // Clear display
        self.clear()?;
        self.display()?;

        Ok(())
    }

    /// Write a command to the display
    fn write_cmd(&mut self, cmd: u8) -> Result<(), I2cError> {
        // Control byte: Co=0, D/C=0 (command)
        self.i2c.write(self.addr, &[0x00, cmd])
    }

    /// Write data to the display
    fn write_data(&mut self, data: &[u8]) -> Result<(), I2cError> {
        // Control byte: Co=0, D/C=1 (data)
        // Split data into chunks of 31 bytes (32 - 1 for control byte)
        const CHUNK_SIZE: usize = 31;
        
        for chunk in data.chunks(CHUNK_SIZE) {
            let mut buffer = [0u8; 32];
            buffer[0] = 0x40; // Control byte for data
            buffer[1..1 + chunk.len()].copy_from_slice(chunk);
            self.i2c.write(self.addr, &buffer[..1 + chunk.len()])?;
        }

        Ok(())
    }

    /// Clear the display buffer
    pub fn clear(&mut self) -> Result<(), I2cError> {
        let buffer_size = (self.size.width() as usize) * (self.size.pages() as usize);
        for i in 0..buffer_size {
            self.buffer[i] = 0;
        }
        Ok(())
    }

    /// Set a pixel in the buffer
    pub fn set_pixel(&mut self, x: u8, y: u8, on: bool) {
        if x >= self.size.width() || y >= self.size.height() {
            return;
        }

        let page = y / 8;
        let bit = y % 8;
        let idx = (page as usize) * (self.size.width() as usize) + (x as usize);

        if on {
            self.buffer[idx] |= 1 << bit;
        } else {
            self.buffer[idx] &= !(1 << bit);
        }
    }

    /// Draw a character at position (x, y)
    /// Returns the width of the drawn character
    pub fn draw_char(&mut self, x: u8, y: u8, ch: char) -> u8 {
        let glyph = match get_font_5x7(ch) {
            Some(g) => g,
            None => return 0,
        };

        for (col, &byte) in glyph.iter().enumerate() {
            for row in 0..7 {
                let pixel_on = (byte & (1 << row)) != 0;
                let px = x + col as u8;
                let py = y + row;
                if px < self.size.width() && py < self.size.height() {
                    self.set_pixel(px, py, pixel_on);
                }
            }
        }

        6 // 5 pixels width + 1 pixel spacing
    }

    /// Draw a string at position (x, y)
    pub fn draw_string(&mut self, x: u8, y: u8, text: &str) {
        let mut cursor_x = x;
        for ch in text.chars() {
            if cursor_x >= self.size.width() {
                break;
            }
            let char_width = self.draw_char(cursor_x, y, ch);
            cursor_x += char_width;
        }
    }

    /// Update the display with the buffer content
    pub fn display(&mut self) -> Result<(), I2cError> {
        // Set column address range
        self.write_cmd(CMD_SET_COLUMN_ADDR)?;
        self.write_cmd(0)?;
        self.write_cmd(self.size.width() - 1)?;

        // Set page address range
        self.write_cmd(CMD_SET_PAGE_ADDR)?;
        self.write_cmd(0)?;
        self.write_cmd(self.size.pages() - 1)?;

        // Write buffer to display
        let buffer_size = (self.size.width() as usize) * (self.size.pages() as usize);
        
        // Copy buffer to avoid borrow conflict
        let mut temp_buffer = [0u8; 1024];
        temp_buffer[..buffer_size].copy_from_slice(&self.buffer[..buffer_size]);
        self.write_data(&temp_buffer[..buffer_size])?;

        Ok(())
    }

    /// Set display contrast (0-255)
    pub fn set_contrast(&mut self, contrast: u8) -> Result<(), I2cError> {
        self.write_cmd(CMD_SET_CONTRAST)?;
        self.write_cmd(contrast)?;
        Ok(())
    }

    /// Turn display on
    pub fn turn_on(&mut self) -> Result<(), I2cError> {
        self.write_cmd(CMD_DISPLAY_ON)
    }

    /// Turn display off
    pub fn turn_off(&mut self) -> Result<(), I2cError> {
        self.write_cmd(CMD_DISPLAY_OFF)
    }

    /// Invert display colors
    pub fn invert(&mut self, inverted: bool) -> Result<(), I2cError> {
        if inverted {
            self.write_cmd(CMD_INVERSE_DISPLAY)
        } else {
            self.write_cmd(CMD_NORMAL_DISPLAY)
        }
    }
}

/// Simple 5x7 font for ASCII characters
/// Returns None for unsupported characters
fn get_font_5x7(ch: char) -> Option<[u8; 5]> {
    match ch {
        ' ' => Some([0x00, 0x00, 0x00, 0x00, 0x00]),
        '!' => Some([0x00, 0x00, 0x5F, 0x00, 0x00]),
        '"' => Some([0x00, 0x07, 0x00, 0x07, 0x00]),
        '#' => Some([0x14, 0x7F, 0x14, 0x7F, 0x14]),
        '$' => Some([0x24, 0x2A, 0x7F, 0x2A, 0x12]),
        '%' => Some([0x23, 0x13, 0x08, 0x64, 0x62]),
        '&' => Some([0x36, 0x49, 0x55, 0x22, 0x50]),
        '\'' => Some([0x00, 0x05, 0x03, 0x00, 0x00]),
        '(' => Some([0x00, 0x1C, 0x22, 0x41, 0x00]),
        ')' => Some([0x00, 0x41, 0x22, 0x1C, 0x00]),
        '*' => Some([0x14, 0x08, 0x3E, 0x08, 0x14]),
        '+' => Some([0x08, 0x08, 0x3E, 0x08, 0x08]),
        ',' => Some([0x00, 0x50, 0x30, 0x00, 0x00]),
        '-' => Some([0x08, 0x08, 0x08, 0x08, 0x08]),
        '.' => Some([0x00, 0x60, 0x60, 0x00, 0x00]),
        '/' => Some([0x20, 0x10, 0x08, 0x04, 0x02]),
        '0' => Some([0x3E, 0x51, 0x49, 0x45, 0x3E]),
        '1' => Some([0x00, 0x42, 0x7F, 0x40, 0x00]),
        '2' => Some([0x42, 0x61, 0x51, 0x49, 0x46]),
        '3' => Some([0x21, 0x41, 0x45, 0x4B, 0x31]),
        '4' => Some([0x18, 0x14, 0x12, 0x7F, 0x10]),
        '5' => Some([0x27, 0x45, 0x45, 0x45, 0x39]),
        '6' => Some([0x3C, 0x4A, 0x49, 0x49, 0x30]),
        '7' => Some([0x01, 0x71, 0x09, 0x05, 0x03]),
        '8' => Some([0x36, 0x49, 0x49, 0x49, 0x36]),
        '9' => Some([0x06, 0x49, 0x49, 0x29, 0x1E]),
        ':' => Some([0x00, 0x36, 0x36, 0x00, 0x00]),
        ';' => Some([0x00, 0x56, 0x36, 0x00, 0x00]),
        '<' => Some([0x08, 0x14, 0x22, 0x41, 0x00]),
        '=' => Some([0x14, 0x14, 0x14, 0x14, 0x14]),
        '>' => Some([0x00, 0x41, 0x22, 0x14, 0x08]),
        '?' => Some([0x02, 0x01, 0x51, 0x09, 0x06]),
        '@' => Some([0x32, 0x49, 0x79, 0x41, 0x3E]),
        'A' => Some([0x7E, 0x11, 0x11, 0x11, 0x7E]),
        'B' => Some([0x7F, 0x49, 0x49, 0x49, 0x36]),
        'C' => Some([0x3E, 0x41, 0x41, 0x41, 0x22]),
        'D' => Some([0x7F, 0x41, 0x41, 0x22, 0x1C]),
        'E' => Some([0x7F, 0x49, 0x49, 0x49, 0x41]),
        'F' => Some([0x7F, 0x09, 0x09, 0x09, 0x01]),
        'G' => Some([0x3E, 0x41, 0x49, 0x49, 0x7A]),
        'H' => Some([0x7F, 0x08, 0x08, 0x08, 0x7F]),
        'I' => Some([0x00, 0x41, 0x7F, 0x41, 0x00]),
        'J' => Some([0x20, 0x40, 0x41, 0x3F, 0x01]),
        'K' => Some([0x7F, 0x08, 0x14, 0x22, 0x41]),
        'L' => Some([0x7F, 0x40, 0x40, 0x40, 0x40]),
        'M' => Some([0x7F, 0x02, 0x0C, 0x02, 0x7F]),
        'N' => Some([0x7F, 0x04, 0x08, 0x10, 0x7F]),
        'O' => Some([0x3E, 0x41, 0x41, 0x41, 0x3E]),
        'P' => Some([0x7F, 0x09, 0x09, 0x09, 0x06]),
        'Q' => Some([0x3E, 0x41, 0x51, 0x21, 0x5E]),
        'R' => Some([0x7F, 0x09, 0x19, 0x29, 0x46]),
        'S' => Some([0x46, 0x49, 0x49, 0x49, 0x31]),
        'T' => Some([0x01, 0x01, 0x7F, 0x01, 0x01]),
        'U' => Some([0x3F, 0x40, 0x40, 0x40, 0x3F]),
        'V' => Some([0x1F, 0x20, 0x40, 0x20, 0x1F]),
        'W' => Some([0x3F, 0x40, 0x38, 0x40, 0x3F]),
        'X' => Some([0x63, 0x14, 0x08, 0x14, 0x63]),
        'Y' => Some([0x07, 0x08, 0x70, 0x08, 0x07]),
        'Z' => Some([0x61, 0x51, 0x49, 0x45, 0x43]),
        '[' => Some([0x00, 0x7F, 0x41, 0x41, 0x00]),
        '\\' => Some([0x02, 0x04, 0x08, 0x10, 0x20]),
        ']' => Some([0x00, 0x41, 0x41, 0x7F, 0x00]),
        '^' => Some([0x04, 0x02, 0x01, 0x02, 0x04]),
        '_' => Some([0x40, 0x40, 0x40, 0x40, 0x40]),
        '`' => Some([0x00, 0x01, 0x02, 0x04, 0x00]),
        'a' => Some([0x20, 0x54, 0x54, 0x54, 0x78]),
        'b' => Some([0x7F, 0x48, 0x44, 0x44, 0x38]),
        'c' => Some([0x38, 0x44, 0x44, 0x44, 0x20]),
        'd' => Some([0x38, 0x44, 0x44, 0x48, 0x7F]),
        'e' => Some([0x38, 0x54, 0x54, 0x54, 0x18]),
        'f' => Some([0x08, 0x7E, 0x09, 0x01, 0x02]),
        'g' => Some([0x0C, 0x52, 0x52, 0x52, 0x3E]),
        'h' => Some([0x7F, 0x08, 0x04, 0x04, 0x78]),
        'i' => Some([0x00, 0x44, 0x7D, 0x40, 0x00]),
        'j' => Some([0x20, 0x40, 0x44, 0x3D, 0x00]),
        'k' => Some([0x7F, 0x10, 0x28, 0x44, 0x00]),
        'l' => Some([0x00, 0x41, 0x7F, 0x40, 0x00]),
        'm' => Some([0x7C, 0x04, 0x18, 0x04, 0x78]),
        'n' => Some([0x7C, 0x08, 0x04, 0x04, 0x78]),
        'o' => Some([0x38, 0x44, 0x44, 0x44, 0x38]),
        'p' => Some([0x7C, 0x14, 0x14, 0x14, 0x08]),
        'q' => Some([0x08, 0x14, 0x14, 0x18, 0x7C]),
        'r' => Some([0x7C, 0x08, 0x04, 0x04, 0x08]),
        's' => Some([0x48, 0x54, 0x54, 0x54, 0x20]),
        't' => Some([0x04, 0x3F, 0x44, 0x40, 0x20]),
        'u' => Some([0x3C, 0x40, 0x40, 0x20, 0x7C]),
        'v' => Some([0x1C, 0x20, 0x40, 0x20, 0x1C]),
        'w' => Some([0x3C, 0x40, 0x30, 0x40, 0x3C]),
        'x' => Some([0x44, 0x28, 0x10, 0x28, 0x44]),
        'y' => Some([0x0C, 0x50, 0x50, 0x50, 0x3C]),
        'z' => Some([0x44, 0x64, 0x54, 0x4C, 0x44]),
        '{' => Some([0x00, 0x08, 0x36, 0x41, 0x00]),
        '|' => Some([0x00, 0x00, 0x7F, 0x00, 0x00]),
        '}' => Some([0x00, 0x41, 0x36, 0x08, 0x00]),
        '~' => Some([0x10, 0x08, 0x08, 0x10, 0x08]),
        _ => None,
    }
}

